from PyAutoWLWW import *
from PIL import Image
import os
import os.path as osp
import SimpleITK as sitk

def getImgFromDcm(dcmPath):
    dcm_msg = sitk.ReadImage(dcmPath)
    return sitk.GetArrayFromImage(dcm_msg)

def saveNpAsDcm(data, savePath):
    sitk.WriteImage(sitk.GetImageFromArray(data), savePath)

if __name__ == '__main__':
    wlww = AutoWLWW()

    # base_path = rf"G:\不同部位数据集对\HDRDataSet\胸部-世纪坛\label\20221020-101410-268_frame52.png"
    # img = np.array(Image.open(base_path), dtype=np.uint16)
    # height, width = img.shape
    # wlww.Initialization(height, width)
    # result = wlww.Process(img, 1)
    # Image.fromarray(img).save("org.png")
    # Image.fromarray(result).save("temp.png")

    base_path = r"C:\Users\15519\Desktop\del"
    save_path = r"C:\Users\15519\Desktop\del\save"
    dcm_file_list = [osp.join(base_path, name) for name in os.listdir(base_path) if name.endswith(".dcm")]
    for dcm_file in dcm_file_list:
        imgs = getImgFromDcm(dcm_file)
        frame_count, height, width = imgs.shape
        wlww.Initialization(height, width)

        result_seq = []
        for img in imgs:
            result_seq.append(wlww.Process(img, 1))
        result_seq = np.array(result_seq)
        base_name = osp.basename(dcm_file)
        saveNpAsDcm(result_seq, osp.join(save_path, base_name))








